package com.gel.aiagent.agent;

import cn.hutool.core.util.StrUtil;
import com.gel.aiagent.agent.model.AgentState;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;

@Data
@Slf4j
public abstract class BaseAgent
{
    //核心属性
    private String name;

    //提示词
    private  String systemPrompt;

    private String nextStepPrompt;

    //执行控制
    private  int maxSteps=10;

    private  int currentStep;


    private AgentState state=AgentState.IDLE;
    //LLM
    private ChatClient chatClient;
    //Memory
    private List<Message> messageList=new ArrayList<>();




    public  String run(String userPrompt){
        if(this.state!=AgentState.IDLE){
            throw  new RuntimeException("Cannot run from this state:"+ this.state);
        }
        if(StrUtil.isBlank(userPrompt)){
            throw  new RuntimeException("Cannot dun agent with empty user prompt");
        }

        //更改状态
        state=AgentState.RUNNING;

        //记录消息上下文
        messageList.add(new UserMessage(userPrompt));
        //保存结果列表
        List<String> results=new ArrayList<>();

        try {
            for(int i=0;i<maxSteps&state!=AgentState.FINISHED;i++){
                int stepNumber=i+1;
                currentStep=stepNumber;
                log.info("Executing step"+stepNumber +"/" + maxSteps);

                //单步执行
                String stepResult=step();
                String result="Step"+stepNumber+": "+stepResult;
                results.add( result);
            }
            if(currentStep>=maxSteps){
                state=AgentState.FINISHED;
                results.add("Terminate: Reached max steps("+maxSteps+")");
            }
            return String.join("\n",results);
        } catch (Exception e) {
            state=AgentState.ERROR;
            log.error("Error executing agent",e);
            return "Executing error:"+e.getMessage();
        }finally {
            this.cleanup();
        }
    }
    /*流失输出*/
    public SseEmitter runStream(String userPrompt){
        SseEmitter sseEmitter = new SseEmitter(300000L);
        CompletableFuture.runAsync(()->{
            try {  if(this.state!=AgentState.IDLE){
                sseEmitter.send("Cannot run from this state:"+ this.state);
                sseEmitter.complete();
                return;
            }
                if(StrUtil.isBlank(userPrompt)){
                    sseEmitter.send("Cannot dun agent with empty user prompt");
                    sseEmitter.complete();
                    return;
                }

                //更改状态
                state=AgentState.RUNNING;

                //记录消息上下文
                messageList.add(new UserMessage(userPrompt));


                try {
                    for(int i=0;i<maxSteps&state!=AgentState.FINISHED;i++){
                        int stepNumber=i+1;
                        currentStep=stepNumber;
                        log.info("Executing step"+stepNumber +"/" + maxSteps);

                        //单步执行
                        String stepResult=step();
                        String result="Step"+stepNumber+": "+stepResult;
                        sseEmitter.send( result);
                    }
                    if(currentStep>=maxSteps){
                        state=AgentState.FINISHED;
                        sseEmitter.send("Terminate: Reached max steps("+maxSteps+")");
                    }
                    sseEmitter.complete();
                } catch (Exception e) {
                    state=AgentState.ERROR;
                    log.error("Error executing agent",e);
                    sseEmitter.send("Executing error:"+e.getMessage());
                    sseEmitter.complete();
                }finally {
                    this.cleanup();
                }
            } catch (Exception e) {
                sseEmitter.completeWithError( e);
            }

        });
        //设置超超时和完成回调
        sseEmitter.onTimeout(() -> {
            this.state=AgentState.ERROR;
            this.cleanup();
            log.warn("SSE connection timeout");
        });
        sseEmitter.onCompletion(() -> {
            if(this.state==AgentState.RUNNING){
                this.state=AgentState.FINISHED;
            }
            this.cleanup();
            log.info("SSE connection completed");
        });
        return sseEmitter;

    }

    private void cleanup() {
    }


    public abstract String step();
}
